import os
from typing import Union
import pandas as pd

from langchain_openai import ChatOpenAI
from langchain.output_parsers import PydanticOutputParser
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

from .python_script_file_model import PythonScriptFileModel
from utils import LOG

class GeneratorChain:
    def __init__(self, model_name: str = "gpt-3.5-turbo", verbose: bool = True, api_key: str = "", base_url: str = "", common_paths: dict = {}):
        self.common_paths = common_paths
        parser = PydanticOutputParser(pydantic_object=PythonScriptFileModel)
        template = (
            f"""
            # Role:
            Python脚本编写专家
            # Goal:
            根据用户提供的[example数据]，编写一个python脚本实现更多类似的fake数据生成。
            # Output Format:
            {parser.get_format_instructions()}
            # Constrains:
            - 只能使用的第三方依赖库：Faker、jsonlines
            - 如果[example数据]的属性有在[example数据属性seed数据规则]做说明，则按照[example数据属性seed数据规则]进行fake数据生成
            - [example数据属性seed数据规则]中的seed_data_seeker是用于查询seed data的外部module
            - 'seed_data_seeker.find'用于返回满足参数条件的seed data数据列表，脚本中需要在返回的数据列表中使用'random.chioce'方法选择某个值作为fake数据
            - [example数据属性seed数据规则]中[seed数据来源]一列值为'含变量字符串'时，则需要替换数据列表中以'[]'进行标识的部分，比如'[timestamp]'则需要替换为时间戳
            - 请直接import seed_data_seeker来加载该模块，而不需要在脚本中进行实现
            - 需要将生成的fake数据保存到'{self.common_paths.get('fake-data-output-files-path')}'目录下的'fake_data_gen.[ext]'文件中，[ext]为文件扩展名（只能使用: jsonl、sql），请根据fake数据类型进行调整
            - 生成新的'fake_data_gen.[ext]'文件前，请先删除旧的'fake_data_gen.[ext]'文件
            - 请确保脚本中使用到的modules都被正确import
            """
        )
        system_message_prompt = SystemMessagePromptTemplate.from_template(template = template, template_format = "jinja2")

        human_template = "# Example Data:\n{example}\n# Seed Data Rules:\n{rules}"
        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)

        # 使用 System 和 Human 角色的提示模板构造 ChatPromptTemplate
        chat_prompt_template = ChatPromptTemplate.from_messages(
            [system_message_prompt, human_message_prompt]
        )
        LOG.debug(f"Chat prompt template: {chat_prompt_template}")

        # 为了翻译结果的稳定性，将 temperature 设置为 0
        chat = ChatOpenAI(openai_api_key=api_key, 
                          model_name=model_name,
                          openai_api_base=base_url, 
                          temperature=0, 
                          verbose=verbose)

        self.chain = chat_prompt_template | chat | parser

    def dataframe_to_markdown_table(self, df):  
        # 定义表头和分隔符
        header = "| property | seed_data_type | seed_data_value |\n| -- | -- | -- |\n"  
        row_template = "| {name} | {source} | {value} |\n"  
        
        # 将DataFrame转换为字符串  
        table_str = header  
        for _, row in df.iterrows():  
            table_str += row_template.format(name=row['property'], source=row['seed_data_type'], value=row['seed_data_value'])  
            
        return table_str

    def run(self, example: str, rules_df: pd.DataFrame) -> Union[str, str, bool]:
        result = None
        try:
            result = self.chain.invoke({
                "example": example,
                "rules": self.dataframe_to_markdown_table(rules_df)
            })
        except Exception as e:
            LOG.error(f"An error occurred during generation: {e}")
            return None, False
        LOG.debug(f'generated python script: {result}')
        generator_file_path = os.path.join(self.common_paths.get('fake-data-generator-path'), result.script_filename)
        with open(generator_file_path, 'w', encoding='utf-8') as file:  
            file.write(result.python_script)
        module_name, _ = os.path.splitext(result.script_filename)
        return module_name, result.data_type, True
